import torch
from PIL import Image
import numpy as np

from model import Model, text_l
from tokenizer import decode, encode

model = Model()
state_dict = torch.load("model.pt")
model.load_state_dict(state_dict)

img = Image.open("../test.png")
img = np.array(img)
img = torch.tensor(img).float()
img = img / 255.0
img = img.permute([2, 1, 0])

idx = torch.zeros([text_l], dtype=torch.long)
idx[0] = encode("\u0002")[0]
pos = 0

while(True):
    logits = model(img, idx)
    y = torch.argmax(logits, -1)
    generated_id = y[pos]
    generated_str = decode([generated_id])
    # print(generated_id, end="")

    if generated_str == "\u0003" or pos == 31:
        break

    print("{}".format(generated_str), end="")
    pos += 1
    idx[1:pos+1] = y[0:pos]

print("")